import struct
from symtable import Class
from typing import List, Optional, Tuple, Any, Iterator
from src.storage.page import Page, PageType
from src.storage.buffer_pool import BufferPoolManager
from src.storage.data_manager import DataManager
from .node import Node
from .sub_array import SubArray
from src.storage.data_item import DataItem
from src.utils.parser import Parser


class InsertRes:
    def __init__(self, new_node: int, new_key: int):
        self.new_node = new_node
        self.new_key = new_key


class BTree:
    """B+树索引
    
    实现B+树索引结构，支持插入、查找、范围查询操作
    使用Node类实现节点操作
    """
    def __init__(self,  data_manager: DataManager, boot_uid:int,data_item:DataItem) -> None:
        self.dm = data_manager  # 数据管理器
        self.boot_uid = boot_uid
        self.boot_data_item = data_item
        self.height = 0
    #     返回一个uid
    @classmethod
    def create(cls,data_manager:DataManager) -> 'BTree':
        raw_root=Node.new_nil_root_raw()
        # boot其实相当于一个哨兵节点，避免了根节点替换导致的树消失
        root_uid=data_manager.insert(0,raw_root)
        # 返回的是uid
        return data_manager.insert(0,Parser.long_to_bytes(root_uid))
        
    @classmethod
    def load(cls,boot_uid:int, data_manager:DataManager) -> 'BTree':
        # boot_data_item处
        boot_data_item=data_manager.read(boot_uid)
        return cls(data_manager, boot_uid, boot_data_item)

    def root_uid(self):
        sa=self.boot_data_item.data()
        # 补0
        return Parser.parse_long(sa.raw[sa.start:sa.start+8])
    # 更新根节点
    def update_root_uid(self,left:int,right:int,rightKey:int) -> None:
        root_raw=Node.new_root_raw(left,right,rightKey)
        new_root_uid=self.dm.insert(0,root_raw)
        self.boot_data_item.before()
        di_raw=self.boot_data_item.data()
        di_raw.raw[di_raw.start:di_raw.start + 8] = Parser.long_to_bytes(new_root_uid)
        self.boot_data_item.after(0)

    def insert(self, key: int, uid: int) -> None:
        """插入键值对
        
        Args:
            key: 索引键
            uid: 数据项UID
            
        Returns:
            操作是否成功
        """
        # 确保根节点存在
        root_uid = self.root_uid()
        
        # 执行插入操作
        insert_result = self._insert(root_uid, key, uid)
        
        # 如果需要创建新的根节点
        if insert_result.new_son != 0:
            # 更新根节点
            self.update_root_uid(root_uid, insert_result.new_son, insert_result.new_key)
            self.height += 1
        
        return True

    def _insert(self, node_uid: int, key: int, uid: int) -> Node.InsertAndSplitRes:
        """递归插入键值对"""
        node = Node.load_node(self, node_uid)
        
        result = Node.InsertAndSplitRes()
        if node.is_leaf():
            # 叶子节点，执行插入并可能分裂
            result = self._insert_and_split(node, key, uid)
        else:
            # 内部节点，找到下一个节点并递归插入
            search_res = node.search_next(key)
            if search_res.uid != 0:
                # 继续在子节点中插入
                insert_result = self._insert(search_res.uid, key, uid)
                
                # 如果子节点分裂，需要在当前节点插入新键
                if insert_result.new_son != 0:
                    result = self._insert_and_split(node, insert_result.new_key, insert_result.new_son)
            else:
                # 需要在兄弟节点中插入
                result.sibling_uid = search_res.sibling_uid
        
        node.release()
        return result

    def _insert_and_split(self, node: Node, key: int, uid: int) -> Node.InsertAndSplitRes:
        """插入键值并在必要时分裂节点"""
        while True:
            result = node.insert_and_split(uid, key)
            
            # 检查是否需要继续在兄弟节点插入
            if result.sibling_uid != 0:
                node.release()
                node = Node.load_node(self, result.sibling_uid)
            else:
                # 返回分裂结果
                return result

    def search(self, key: int) -> List[int]:
        """查找指定键值"""
        return self.search_range(key, key)

    def search_range(self, left_key: int, right_key: int) -> List[int]:
        """范围查询"""

        # 查找包含左键的叶子节点
        leaf_uid = self._search_leaf(self.root_uid(), left_key)
        
        # 收集结果
        results = []
        while leaf_uid != 0:
            node = Node.load_node(self, leaf_uid)
            
            # 在叶子节点中搜索范围
            leaf_results = node.leaf_search_range(left_key, right_key)
            node.release()

            # 检查是否需要继续到下一个叶子节点
            next_leaf_uid = leaf_results.sibling_uid
            results.extend(leaf_results.uids)

            if next_leaf_uid == 0:
                break
            else:
                leaf_uid = next_leaf_uid
        
        return results

    def _search_leaf(self, node_uid: int, key: int) -> int:
        """查找包含指定键的叶子节点"""
        node = Node.load_node(self, node_uid)
        
        if node.is_leaf():
            uid = node_uid
            node.release()
            return uid
        else:
            search_res = node.search_next(key)
            node.release()
            
            if search_res.uid != 0:
                return self._search_leaf(search_res.uid, key)
            else:
                return 0  # 未找到

    def close(self) -> None:
        """关闭B+树"""
        # 释放资源
        self.boot_data_item.release()

